import torch
import torch.nn as nn

# 定义一个简单的神经网络模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        # 定义一个输入层到隐藏层的全连接层
        self.fc1 = nn.Linear(2, 2) # 输入2个特征，输出2个特征
        # 定义一个隐藏层到输出层的全连接层
        self.fc2 = nn.Linear(2, 1) # 输入 2 个特征，输出 1 个预测值
    def forward(self, x):
        # 前向传播过程
        x = torch.relu(self.fc1(x)) # 使用ReLU激活函数
        x = self.fc2(x) # 输出层
        return x
    
# 创建模型实例
model = SimpleNN()
print(model)

# 损失函数
criterion = nn.MSELoss()

# 使用 SGD 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练数据示例
X = torch.randn(10, 2) # 10 个样本，每个样本有 2 个特征
Y = torch.randn(10, 1) # 10 个目标标签

# 训练过程
for epoch in range(100): # 训练100轮
    model.train() # 设置模型为训练模式
    optimizer.zero_grad() # 清除梯度
    output = model(X) # 前向传播
    loss = criterion(output, Y) # 计算损失
    loss.backward() # 反向传播
    optimizer.step() # 更新权重

    if (epoch + 1) % 10 == 0: # 每 10 轮输出一次损失
        print(f'Epoch[{epoch + 1}/100], Loss:{loss.item():.4f}')

